package ai.zixing.mashibing.new_class.class06;

import java.util.Comparator;
import java.util.PriorityQueue;

// https://leetcode.cn/problems/merge-k-sorted-lists/
// 设有M条链表，总节点个数为N
// 小根堆大小不会超过 M，小根堆的所有操作时 O(lgM)
// N 个节点，出来进去总共 O(N * lgM)
public class Code01_MergeKSortedLists {

    public static class ListNode {
        public int val;
        public ListNode next;
    }

    public static class ListNodeComparator implements Comparator<ListNode> {
        @Override
        public int compare(ListNode o1, ListNode o2) {
            return o1.val -o2.val;
        }
    }

    public static ListNode mergeKLists(ListNode[] lists) {
        if (lists == null) {
            return null;
        }
        PriorityQueue<ListNode> heap = new PriorityQueue<>(new ListNodeComparator());
        // 将所有链表的头节点放入小根堆
        for (ListNode list : lists) {
            if (list != null) {
                heap.offer(list);
            }
        }
        if (heap.isEmpty()) {
            return null;
        }
        // 取出小根据第一个元素
        ListNode head = heap.poll();
        ListNode pre = head;
        // 初始化输出链表
        if (head.next != null) {
            heap.offer(pre.next);
        }
        while (!heap.isEmpty()) {
            ListNode cur = heap.poll();
            pre.next = cur;
            pre = cur;
            if (cur.next != null) {
                heap.offer(cur.next);
            }
        }
        return head;
    }
}
